from torch.utils.data import DataLoader
BASE_PATH = '/public_data/zzh/ts_exp'
mapping = {
    'electricity':{
        'root_path':f'{BASE_PATH}/dataset/electricity',
        'file_path':'electricity.csv',
        'features':'M',
        'target': 'OT',
        'timeenc': 0,
        'freq': 'h'
    },
    'PEMS_03':{
        'root_path':f'{BASE_PATH}/dataset/PEMS',
        'file_path':'PEMS03.npz',
        'features':'M',
        'target': 'OT',
        'timeenc': 0,
        'freq': 'h'
    },
    'PEMS_04':{
        'root_path':f'{BASE_PATH}/dataset/PEMS',
        'file_path':'PEMS04.npz',
        'features':'M',
        'target': 'OT',
        'timeenc': 0,
        'freq': 'h'
    },
    'PEMS_07':{
        'root_path':f'{BASE_PATH}/dataset/PEMS',
        'file_path':'PEMS07.npz',
        'features':'M',
        'target': 'OT',
        'timeenc': 0,
        'freq': 'h'
    },
    'PEMS_08':{
        'root_path':f'{BASE_PATH}/dataset/PEMS',
        'file_path':'PEMS08.npz',
        'features':'M',
        'target': 'OT',
        'timeenc': 0,
        'freq': 'h'
    },
    'weather': {
        'root_path': f'{BASE_PATH}/dataset/weather',
        'file_path': 'weather.csv',
        'features': 'M',
        'target': 'OT',
        'timeenc': 0,
        'freq': 'h'
    },
    'exchange': {
        'root_path': f'{BASE_PATH}/dataset/exchange_rate',
        'file_path': 'exchange_rate.csv',
        'features': 'M',
        'target': 'OT',
        'timeenc': 0,
        'freq': 'h'
    },
    'traffic':{
        'root_path':f'{BASE_PATH}/dataset/traffic',
        'file_path':'traffic.csv',
        'features':'M',
        'target': 'OT',
        'timeenc': 0,
        'freq': 'h'
    },
    'Solar':{
      'file_path':'solar_AL.txt',
      'root_path': f'{BASE_PATH}/dataset/Solar',
      'freq': 'h',
      'target':'OT',
      'features': 'M',
      'timeenc':1,
    },
    'ETTh1':{
        'root_path': f'{BASE_PATH}/dataset/ETT-small',
        'file_path': 'ETTh1.csv',
        'target':'OT',
        'features': 'M',
        'timeenc':0,
        'freq':'h'
    },
    'ETTh2': {
        'root_path': f'{BASE_PATH}/dataset/ETT-small',
        'file_path': 'ETTh2.csv',
        'target':'OT',
        'features': 'M',
        'timeenc': 0,
        'freq': 'h'

    },
    'ETTm1': {
        'root_path': f'{BASE_PATH}/dataset/ETT-small',
        'file_path': 'ETTm1.csv',
        'target':'OT',
        'features': 'M',
        'timeenc': 0,
        'freq': 'm'
    },
    'ETTm2': {
        'root_path': f'{BASE_PATH}/dataset/ETT-small',
        'file_path': 'ETTm2.csv',
        'features': 'M',
        'target':'OT',
        'timeenc': 0,
        'freq': 'm'
    },

}

def get_loader(Data, batch_size, data, look_back, pre_win):
    trainset = Data(
        root_path=mapping[data]['root_path'],
        data_path=mapping[data]['file_path'],
        flag='train',
        size=[look_back, 0, pre_win],
        features=mapping[data]['features'],
        target=mapping[data]['target'],
        timeenc=mapping[data].get('timeenc', 0),
        freq=mapping[data]['freq'],
    )
    testset = Data(
        root_path=mapping[data]['root_path'],
        data_path=mapping[data]['file_path'],
        flag='test',
        size=[look_back, 0, pre_win],
        features=mapping[data]['features'],
        target=mapping[data]['target'],
        timeenc=mapping[data].get('timeenc', 0),
        freq=mapping[data]['freq'],
    )
    valiset = Data(
        root_path=mapping[data]['root_path'],
        data_path=mapping[data]['file_path'],
        flag='val',
        size=[look_back, 0, pre_win],
        features=mapping[data]['features'],
        target=mapping[data]['target'],
        timeenc=mapping[data].get('timeenc', 0),
        freq=mapping[data]['freq'],
    )
    train_loader = DataLoader(
        trainset,
        batch_size=batch_size,
        shuffle=False,
        num_workers=8,
        drop_last=False)
    test_loader = DataLoader(
        testset,
        batch_size=1,
        shuffle=False,
        num_workers=8,
        drop_last=True)
    vali_loader = DataLoader(
        valiset,
        batch_size=batch_size,
        shuffle=False,
        num_workers=8,
        drop_last=False)
    return test_loader, train_loader, vali_loader
